from BasicLibraries import *
import regression as rg
from matplotlib import rc

plt.rcParams['xtick.major.pad']='8'
rc('text', usetex=True)
rc('font', size=24)
rc('legend', fontsize=24)
rc('text.latex', preamble=r'\usepackage{cmbright}')
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})

color = ['#e95057', '#0080c4', '#00a989', '#464646']


#%% Effect of diversification
def estimate_effect_of_BRD_in_time(dt_ws, dummies, YFE, sdgs_options_name, diversification_type, standardise_coefficients=True):
    print(sdgs_options_name, diversification_type)
    tab = pd.DataFrame()
    dep_var = 'DirectControl_fut_tot'
    res_bar = pd.DataFrame()
    count=0
    list_of_models = dict()
    list_of_data = dict()
    for Y in range(2013,2020):
        #==
        Z = dt_ws[dt_ws.rfyear <= Y]
        new_dummies = dummies.copy()
        if YFE == True:
            for NY in range(Y+1,2020):
                new_dummies.remove(NY)
        res, resNUM, full_table, x, models, baseline = rg.make_impact_emission_regression(Z, new_dummies,  emission_var = dep_var,
                                                    years_FE=YFE, standardise_coefficients = standardise_coefficients)
        res.index = [Y]
        resNUM[3] = (resNUM[2]-resNUM[1])/2
        resNUM.columns = ['Effect', 'LB', 'UB', 'Uncertainty']
        resNUM.index = res.columns
        resNUM = resNUM.rename(index = {'Risk mitigation': 'Risk Mit.', 'Stakeholders engagement': 'Stak.Eng.',
                                        'Diversification': 'Response diversity'})
        resNUM['Year'] = [Y]*len(resNUM)
        res_bar = pd.concat((res_bar, resNUM))
        count+=1
        tab = pd.concat((tab, res))
        list_of_models[Y]  = models
        list_of_data[Y] = x
    
    print(tab.to_latex())
    shifts = [-0.1,-0.05,0.05,0.1]
    count = 0
    ax = plt.subplot()
    for m in resNUM.index.unique():
        tmp1 = res_bar.loc[m].set_index('Year')
        tmp1.index +=shifts[count]
        tmp1 = tmp1.rename(columns = {'Effect': m})
        tmp1[m].plot(yerr = tmp1['Uncertainty'], ls = '', 
                            marker = 'o', ms = 15, capsize = 4, c = color[count], ax = ax)
        count+=1
    plt.axhline(y = 0, ls ='-.', lw = 0.5, c = 'k')
    if standardise_coefficients:
        plt.ylabel('Impact on emissions')
    else:
        plt.ylabel('Impact on emissions\nUn-standardised coefficients')       
    plt.legend(loc = 'center', bbox_to_anchor = (0.5,1.1), ncol=4)
    plt.xlabel('')
    
    if standardise_coefficients==True:
        ax.text(-0.1, 1., 'B', transform=ax.transAxes,
                 fontsize=24, fontweight='bold', va='top', ha='right')
        plt.savefig('Figures/Figure4B.pdf', dpi = 100, bbox_inches = 'tight')
    else:
        plt.savefig('Figures/Figure4B_SI.pdf', dpi = 100, bbox_inches = 'tight')

    return res_bar, list_of_models, list_of_data, tab